import json
import os

import cv2


class PixelClicker:
    """
    Manual quest generator for the robot. Given a message like:

    quest = {
        "left_arm": {"action": "raise_to", "pixels": [0.6]},
        "right_arm": {
            "action": "pick_lift",
            "pixels": [np.array([0.408, 0.2, 0.045]), 0.45],
            "scale": 0.0,
        },
    }

    the "pixels" will be transformed to "args" that is used for the robot.
    """

    def __init__(
        self,
    ):
        self.img_path = ".cache/imgs/current_view.jpg"
        self.output_path = ".cache/txts/message_response.json"
        self.original_image = cv2.imread(self.img_path)
        self.window_name = "Clothes folding"
        self.scale = 1.0
        self.offset_x = 0
        self.offset_y = 0
        self.dragging = False
        self.last_x, self.last_y = 0, 0
        self.pixels = {
            "left": {"action": "", "pixels": [(None, None), (None, None)], "height": 0.0},
            "right": {"action": "", "pixels": [(None, None), (None, None)], "height": 0.0},
        }
        self.camera_xyz = []
        self.gripper_xyz = []
        self.colors = [(0, 255, 255), (0, 0, 255), (255, 255, 0), (255, 0, 255)]
        #    "left_pick": (0, 255, 255),  # 黄色
        #     "left_place": (0, 0, 255),  # 红色
        #     "right_pick": (255, 255, 0),  # 青色
        #     "right_place": (255, 0, 255),  # 紫色
        # }

        self.click_order = ["start", "end"]

        self.skills = [
            "pick_place",
            "raise_to",
            "pick_lift",
            "place_down",
            "press_hold",
            "line_move",
            "swing_forward",
        ]
        self.skill_id = 0
        self.selected_arm = "left"
        self.action_height = 0.0

    def display_coordinates(self, event, x, y, flags, param):
        """鼠标事件回调函数"""
        # 将显示坐标转换为原始图像坐标
        original_x = int((x + self.offset_x) / self.scale)
        original_y = int((y + self.offset_y) / self.scale)

        # 确保坐标在有效范围内
        if (
            original_x < 0
            or original_x >= self.original_image.shape[1]
            or original_y < 0
            or original_y >= self.original_image.shape[0]
        ):
            return

        if event == cv2.EVENT_LBUTTONDOWN:
            # Click: add a task+pixel for the arm
            # click to save the skill
            pixel = [original_x, original_y]
            for i, p in enumerate(self.pixels[self.selected_arm]["pixels"]):
                if p == (None, None):
                    self.pixels[self.selected_arm]["pixels"][i] = pixel
                    self.pixels[self.selected_arm]["action"] = self.skills[self.skill_id]
                    break

                print(
                    f"Add {self.selected_arm} {self.skills[self.skill_id]} pixel=({original_x}, {original_y}), "
                    f"action_scale={self.action_height:.3f}, color={self.original_image[original_y, original_x]}"
                )

        elif event == cv2.EVENT_RBUTTONDOWN:
            if self.pixels[self.selected_arm]["pixels"]:
                for i, p in enumerate(self.pixels[self.selected_arm]["pixels"][::-1]):
                    if p != (None, None):
                        removed = self.pixels[self.selected_arm]["pixels"][1 - i]
                        self.pixels[self.selected_arm]["pixels"][1 - i] = (None, None)
                        break
                print(
                    f"del {self.selected_arm} arm: {self.pixels[self.selected_arm]['action']} "
                    f"pixel=({removed[0]}, {removed[1]})"
                )

        elif event == cv2.EVENT_MBUTTONDOWN:
            # 中键点击：开始拖动
            self.dragging = True
            self.last_x, self.last_y = x, y

        elif flags & cv2.EVENT_FLAG_SHIFTKEY and event == cv2.EVENT_MOUSEMOVE:
            self.action_height += 0.001
            self.action_height = max(0.0, min(1.0, self.action_height))
            print(f"Action scale: {self.action_height:.3f}")
            self.pixels[self.selected_arm]["height"] = round(self.action_height, 3)
        elif flags & cv2.EVENT_FLAG_CTRLKEY and event == cv2.EVENT_MOUSEMOVE:
            self.action_height -= 0.001
            self.action_height = max(0.0, min(1.0, self.action_height))
            print(f"Action scale: {self.action_height:.3f}")
            self.pixels[self.selected_arm]["height"] = round(self.action_height, 3)
        # 显示当前鼠标位置的坐标
        cv2.displayOverlay(
            self.window_name,
            f"Pixel: ({original_x}, {original_y}) {self.selected_arm}: {self.skills[self.skill_id]},"
            f"height={self.pixels[self.selected_arm]['height']:.3f}",
            100,
        )

    def run(self):
        """运行标定工具"""

        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.setMouseCallback(self.window_name, self.display_coordinates)

        while True:
            # 创建当前视图的副本
            display = self.original_image.copy()  # cv2.imread("imgs/current_view.jpg")
            height, width = display.shape[:2]
            # 绘制所有标记点

            for arm in ["left", "right"]:
                start_color_id = 0 if arm == "left" else 2
                tag = f"{arm} {self.pixels[arm]['action']}"
                for i, (px, py) in enumerate(self.pixels[arm]["pixels"]):
                    if px is not None:
                        color = self.colors[start_color_id + i]
                        cv2.circle(display, (px, py), 2, color, -1)
                        cv2.putText(display, tag, (px + 10, py - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

            # display all info
            status = f"{self.selected_arm}: {self.skills[self.skill_id]}, height={self.pixels[self.selected_arm]['height']:.3f}"
            cv2.putText(
                display,
                status,
                (20, height - 10),
                cv2.FONT_HERSHEY_SIMPLEX,
                1,
                (40 * int(self.skill_id), 0, 255),
                2,
            )

            # 应用缩放和偏移
            h, w = display.shape[:2]
            view_x1 = max(0, int(self.offset_x / self.scale))
            view_y1 = max(0, int(self.offset_y / self.scale))
            view_x2 = min(w, int((self.offset_x + w) / self.scale))
            view_y2 = min(h, int((self.offset_y + h) / self.scale))

            view = display[view_y1:view_y2, view_x1:view_x2]
            view = cv2.resize(view, (w, h))

            # 显示图像
            cv2.imshow(self.window_name, view)

            # 处理键盘事件
            key = cv2.waitKey(1) & 0xFF

            if key == 27:  # ESC键退出
                # if os.path.exists(self.output_path):
                #     os.remove(self.output_path)
                break
            elif key == ord("+") or key == ord("="):  # 放大
                self.scale *= 1.1
            elif key == ord("-"):  # 缩小
                self.scale /= 1.1
            elif key == ord("r"):  # reload the image
                self.original_image = cv2.imread(self.img_path)
                # remove all the points
                self.pixels = {
                    "left": {"action": "", "pixels": [(None, None), (None, None)], "height": self.action_height},
                    "right": {"action": "", "pixels": [(None, None), (None, None)], "height": self.action_height},
                }
                self.scale = 1.0
                self.offset_x = 0
                self.offset_y = 0
            elif key == ord("s"):  # Save the points
                cv2.imwrite(".cache/imgs/annotated_current_view.jpg", display)
                print(f"Saved annotated image to: .cache/imgs/annotated_current_view.jpg")
                # send the message to the robot: use the same as LLMs
                with open(self.output_path, "w") as f:
                    json.dump(self.pixels, f)
                print(f"Send txt messages to {self.output_path} ")

            elif key in [ord(str(_)) for _ in range(1, len(self.skills) + 1)]:
                """
                select the skill for the arms
                """
                self.skill_id = key - ord("1")
                print(f"{self.selected_arm} Arm Selected skill: {self.skills[self.skill_id]}")

            elif key == ord("a"):
                """
                switch the arm
                """
                self.selected_arm = "left" if self.selected_arm == "right" else "right"
                print(f"Switch to {self.selected_arm}")

        cv2.destroyAllWindows()


if __name__ == "__main__":

    human_pixel_api = PixelClicker()
    human_pixel_api.run()
